import os
import joblib
import numpy as np
import pandas as pd
import tensorflow as tf
from datetime import datetime

# 设置TensorFlow日志级别
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'  # 隐藏TensorFlow警告

# 模型路径 (修复相对路径问题)
import sys
script_dir = os.path.dirname(os.path.abspath(__file__))
project_root = os.path.dirname(script_dir)
MODEL_DIR = os.path.join(project_root, 'model', 'saved_models')
LSTM_MODEL_PATH = os.path.join(MODEL_DIR, 'lstm_model.h5')
LSTM_SCALERS_PATH = os.path.join(MODEL_DIR, 'lstm_scalers.joblib')

def load_model():
    """加载LSTM模型和缩放器"""
    model = tf.keras.models.load_model(LSTM_MODEL_PATH)
    scalers = joblib.load(LSTM_SCALERS_PATH)
    print(f"LSTM模型已加载")
    return model, scalers

def process_data(data):
    """处理输入数据，生成特征"""
    # 将输入数据转换为DataFrame
    if isinstance(data, str):
        # 如果输入是CSV格式的字符串，解析它
        columns = ['Date', 'Time', 'S1_Temp', 'S2_Temp', 'S3_Temp', 'S4_Temp',
                   'S1_Light', 'S2_Light', 'S3_Light', 'S4_Light',
                   'S1_Sound', 'S2_Sound', 'S3_Sound', 'S4_Sound',
                   'S5_CO2', 'S5_CO2_Slope', 'S6_PIR', 'S7_PIR']
        values = data.strip().split(',')
        
        # 确保值的数量与列数匹配
        if len(values) != len(columns):
            print(f"警告: 输入数据有 {len(values)} 列，但期望 {len(columns)} 列。")
            if len(values) < len(columns):
                # 如果值不够，添加默认值
                values.extend(['0'] * (len(columns) - len(values)))
            else:
                # 如果值太多，截断
                values = values[:len(columns)]
        
        df = pd.DataFrame([values], columns=columns)
        
        # 转换数据类型
        for col in df.columns:
            if col in ['Date', 'Time']:
                continue
            try:
                df[col] = pd.to_numeric(df[col])
            except:
                pass
    else:
        # 如果输入是字典，直接转换为DataFrame
        df = pd.DataFrame([data])
    
    # 添加DateTime列
    if 'Date' in df.columns and 'Time' in df.columns:
        df['DateTime'] = pd.to_datetime(df['Date'] + ' ' + df['Time'], dayfirst=True)
    
    # 添加时间特征
    if 'DateTime' in df.columns:
        df['Hour'] = df['DateTime'].dt.hour
        df['Minute'] = df['DateTime'].dt.minute
        df['DayOfWeek'] = df['DateTime'].dt.dayofweek
        df['IsWeekend'] = (df['DateTime'].dt.dayofweek >= 5).astype(int)
        df['TimeOfDay'] = df['Hour'].apply(lambda x: 0 if 5 <= x < 12 else 1 if 12 <= x < 18 else 2 if 18 <= x < 22 else 3)
        df['IsWorkingHour'] = (((df['Hour'] >= 9) & (df['Hour'] < 18)) & (df['DayOfWeek'] < 5)).astype(int)
    
    # 提取基本传感器数据
    temp_cols = [col for col in df.columns if 'Temp' in col and col.startswith('S')]
    light_cols = [col for col in df.columns if 'Light' in col and col.startswith('S')]
    sound_cols = [col for col in df.columns if 'Sound' in col and col.startswith('S')]
    pir_cols = [col for col in df.columns if 'PIR' in col and col.startswith('S')]
    
    # 计算统计特征
    # 温度特征
    df['Temp_Mean'] = df[temp_cols].mean(axis=1)
    df['Temp_Var'] = df[temp_cols].var(axis=1)
    df['Temp_Max'] = df[temp_cols].max(axis=1)
    df['Temp_Min'] = df[temp_cols].min(axis=1)
    df['Temp_Range'] = df['Temp_Max'] - df['Temp_Min']
    
    # 光线特征
    df['Light_Mean'] = df[light_cols].mean(axis=1)
    df['Light_Var'] = df[light_cols].var(axis=1)
    df['Light_Max'] = df[light_cols].max(axis=1)
    df['Light_Min'] = df[light_cols].min(axis=1)
    
    # 声音特征
    df['Sound_Mean'] = df[sound_cols].mean(axis=1)
    df['Sound_Var'] = df[sound_cols].var(axis=1)
    df['Sound_Max'] = df[sound_cols].max(axis=1)
    
    # PIR特征
    df['PIR_Sum'] = df[pir_cols].sum(axis=1)
    df['PIR_Any'] = (df['PIR_Sum'] > 0).astype(int)
    
    # 创建特征交互项
    df['Temp_Mean_x_Light_Mean'] = df['Temp_Mean'] * df['Light_Mean']
    df['Temp_Mean_x_Sound_Mean'] = df['Temp_Mean'] * df['Sound_Mean']
    df['Temp_Mean_x_S5_CO2'] = df['Temp_Mean'] * df['S5_CO2']
    df['Light_Mean_x_Sound_Mean'] = df['Light_Mean'] * df['Sound_Mean']
    df['Light_Mean_x_S5_CO2'] = df['Light_Mean'] * df['S5_CO2']
    df['Sound_Mean_x_S5_CO2'] = df['Sound_Mean'] * df['S5_CO2']
    
    # 删除不需要的列
    drop_cols = ['Date', 'Time', 'DateTime']
    df = df.drop([col for col in drop_cols if col in df.columns], axis=1)
    
    return df

def predict(data):
    """使用LSTM模型预测房间占用人数"""
    # 加载模型和缩放器
    model, scalers = load_model()
    
    # 处理数据
    processed_data = process_data(data)
    
    # 获取缩放器和序列长度
    important_scaler = scalers['important_scaler']
    other_scaler = scalers['other_scaler']
    y_scaler = scalers['y_scaler']
    important_indices = scalers['important_indices']
    other_indices = scalers['other_indices']
    sequence_length = scalers['sequence_length']
    
    # 提取特征
    X = processed_data.values
    
    # 安全处理：确保特征数量足够
    feature_count = max(max(important_indices + other_indices) + 1 if important_indices or other_indices else 0, X.shape[1])
    if feature_count > X.shape[1]:
        X_safe = np.zeros((X.shape[0], feature_count))
        X_safe[:, :X.shape[1]] = X
        X = X_safe
    
    # 分别缩放重要特征和其他特征
    if len(important_indices) > 0:
        X_important = important_scaler.transform(X[:, important_indices])
    else:
        X_important = np.array([]).reshape(X.shape[0], 0)
    
    if len(other_indices) > 0:
        X_other = other_scaler.transform(X[:, other_indices])
    else:
        X_other = np.array([]).reshape(X.shape[0], 0)
    
    # 合并缩放后的特征
    X_scaled = np.hstack([X_important, X_other]) if X_important.size > 0 and X_other.size > 0 else (X_important if X_important.size > 0 else X_other)
    
    # 创建序列
    X_sequence = np.tile(X_scaled, (1, sequence_length, 1))
    X_sequence = X_sequence.reshape(1, sequence_length, X_scaled.shape[1])
    
    # 预测
    y_pred_scaled = model.predict(X_sequence, verbose=0)
    
    # 反向转换
    y_pred = y_scaler.inverse_transform(y_pred_scaled)
    prediction = y_pred[0][0]
    
    # 四舍五入到最接近的整数
    rounded_prediction = round(float(prediction))
    
    return {
        'model': 'lstm',
        'prediction': float(prediction),
        'rounded_prediction': rounded_prediction
    }

def main():
    """主函数"""
    print("=== LSTM模型房间占用人数预测 ===")
    print("请输入传感器数据 (CSV格式):")
    print("格式: Date,Time,S1_Temp,S2_Temp,S3_Temp,S4_Temp,S1_Light,S2_Light,S3_Light,S4_Light,S1_Sound,S2_Sound,S3_Sound,S4_Sound,S5_CO2,S5_CO2_Slope,S6_PIR,S7_PIR")
    print("例如: 22-12-2017,10:53:46,25,24.81,24.56,25.5,122,35,56,43,0.09,0.05,0.06,0.13,390,-0.1,0,0")
    
    data_input = input("\n输入数据: ")
    if not data_input:
        # 使用示例数据
        data_input = "22-12-2017,10:53:46,25,24.81,24.56,25.5,122,35,56,43,0.09,0.05,0.06,0.13,390,-0.1,0,0"
    
    print("\n预测中...")
    result = predict(data_input)
    
    print("\n预测结果:")
    print(f"预测人数: {result['prediction']}")
    print(f"四舍五入后的人数: {result['rounded_prediction']}")

if __name__ == "__main__":
    main()